import h5py
import json
import argparse
import os
from shutil import copyfile
import robosuite
import xml.etree.ElementTree as ET
from tqdm import tqdm
import numpy as np


def get_eef_panda_site(prefix):
    """
    Helper function for new element added in robosuite v1.5
    """

    eef_panda_site = """<site name="right_center" pos="0 0 0" size="0.01" rgba="1 0.3 0.3 1" group="2"/> """

    tree = ET.ElementTree(ET.fromstring(eef_panda_site))
    root = tree.getroot()
    for elem in root.iter():
        elem.attrib["name"] = "{}_{}".format(prefix, elem.get("name"))

    return tree


def update_gripper_name(gripper_name: str):
    tokens = gripper_name.split("_")
    tokens.insert(1, "right")
    return "_".join(tokens)


def convert_xml(xml_str):
    """
    Postprocess xml string generated by robosuite to be compatible with robosuite v1.5
    This script should not the xml string if it was already generated using robosuite v1.5
    Args:
        xml_str (str): xml string to process (from robosuite v1.4.1)
    """

    root = ET.fromstring(xml_str)
    worldbody = root.find("worldbody")
    bodies_to_process = []

    for body in worldbody.iter("body"):
        body_name = body.get("name")
        if "link0" in body_name:
            if body not in bodies_to_process:
                bodies_to_process.append(body)

    # add elements introduced in v1.5
    for body in bodies_to_process:
        prefix = body.get("name").split("_")[0]
        eef_tree = get_eef_panda_site(prefix)
        body.append(eef_tree.getroot())

    assets = root.find("asset")

    # change path from mounts to bases
    for asset in assets.iter("mesh"):
        path = asset.get("file")
        assert (
            "robosuite/models/assets" in path or "robocasa/models/assets" in path
        ), print(path)
        if "mounts" in path:
            assert not "robot" in path
            new_path = path.replace(
                "robosuite/models/assets/mounts", "robosuite/models/assets/bases"
            )
            asset.set("file", new_path)

    # update naming prefixes
    for elem in root.iter():
        for attr, value in elem.attrib.items():
            if value.startswith("mount"):
                elem.set(attr, value.replace("mount", "fixed_base", 1))
            elif value.startswith("gripper"):
                elem.set(attr, update_gripper_name(value))
